import os
import pandas as pd
import numpy as np
import warnings
from typing import Optional, Dict, List, Tuple

def merge_processed_indicator_files(
    processed_dir: str, 
    output_file: str, 
    year_range: Optional[List[int]] = None
) -> bool:
    """
    合并处理过的指标文件到一个Excel文件中。
    
    参数:
        processed_dir: 处理好的指标文件目录
        output_file: 输出合并后的Excel文件路径
        year_range: 要合并的年份范围，例如[2009, 2023]。如果为None，则使用默认范围(2009-2023)
    
    返回:
        是否成功合并的布尔值
    """
    # 设置默认年份范围
    if year_range is None:
        year_range = list(range(2009, 2024))  # 2009到2023年
    
    # 验证输入参数
    if not isinstance(year_range, list):
        raise TypeError("year_range必须是年份列表")
    
    # 检查并准备输出目录
    output_dir = os.path.dirname(output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)
    
    # 验证输入目录
    if not os.path.isdir(processed_dir):
        raise ValueError(f"指定的目录不存在或不是目录: {processed_dir}")
    
    # 收集所有符合条件的文件
    file_pattern = "_AfterProcess.xlsx"
    indicator_files = []
    
    for file in os.listdir(processed_dir):
        if file.endswith(file_pattern):
            indicator_files.append(os.path.join(processed_dir, file))
    
    # 检查是否有文件可处理
    if not indicator_files:
        warnings.warn(f"在目录 {processed_dir} 中未找到匹配文件模式 '_AfterProcess.xlsx' 的文件")
        try:
            pd.DataFrame(columns=['国名En', '国名Ch', 'Year'] + year_range).to_excel(output_file, index=False)
            print(f"创建了空的输出文件: {output_file}")
            return True
        except Exception as e:
            print(f"警告: 创建空文件时出错 - {str(e)}")
            return False
    
    print(f"发现 {len(indicator_files)} 个匹配的数据文件...")
    
    try:
        # 使用字典存储每个文件的年份列位置信息
        file_year_columns: Dict[str, Dict[str, int]] = {}
        
        # 字典存储合并数据
        merged_data: Dict[str, Dict] = {}
        
        # 处理每个文件
        for idx, file_path in enumerate(indicator_files):
            filename = os.path.basename(file_path)
            print(f" > 处理文件 {idx+1}/{len(indicator_files)}: {filename}")
            
            try:
                # 读取Excel文件，跳过空行和无效字符
                df = pd.read_excel(file_path, na_filter=True, keep_default_na=True)
                
                # 将 '..' 替换为 NaN
                if '..' in df.values:
                    df.replace('..', np.nan, inplace=True)
                    print(f"  - 发现无效值 '..' 并已替换为 NaN")
                
                # 获取列名中的'Year'字段
                year_col = next((col for col in df.columns if 'year' in str(col).lower()), None)
                
                if year_col is None and 'Year' in df.columns:
                    year_col = 'Year'
                elif year_col is None:
                    print(f"  - 警告: 未找到年份列，请在数据中添加'Year'列或包含'year'关键词的列")
                    continue
                
                # 获取国家列
                country_cols = [col for col in df.columns if 'country' in str(col).lower()]
                if not country_cols:
                    print(f"  - 警告: 未找到国家列，请在数据中添加'country'关键词的列")
                    continue
                    
                country_col_en = country_cols[0]
                if 'chinese' in str(country_col_en).lower() and not 'en' in str(country_col_en).lower():
                    # 可能是中文国家列，尝试找英文列
                    country_col_en_pair = next((col for col in df.columns if 'en' in str(col).lower()), None)
                    if country_col_en_pair:
                        country_col_en = country_col_en_pair
                        print(f"  - 使用英文列 '{country_col_en}' 代替中文列 '{country_cols[0]}'")
                
                # 检查是否所有年份都在数据中
                unique_years = sorted(df[year_col].dropna().astype(int).unique())
                missing_years = [str(y) for y in year_range if y not in unique_years]
                if missing_years:
                    print(f"  - 发现缺失年份: {', '.join(missing_years)}")
                
                # 记录当前文件的年份列信息
                file_year_columns[filename] = {'year_col': year_col, 'year_range': year_range}
                
                # 获取指标列（除固定列外的其他列）
                other_cols = [col for col in df.columns 
                            if col not in [year_col, country_col_en, '国名Ch', 'Year'] + [np.nan] * (len(df.columns) - len(df.columns))]
                if not other_cols and '国名Ch' not in df.columns:
                    print(f"  - 警告: 未找到任何指标列，跳过文件 {filename}")
                    continue
                
                # 对于每个指标列，创建一级索引
                for indicator_col in other_cols:
                    for year_val in year_range:
                        # 只处理年份列存在且包含该年份的数据
                        if year_val in unique_years:
                            year_value = df.loc[df[year_col] == year_val, indicator_col].iloc[0] if any(df[year_col] == year_val) else np.nan
                        else:
                            # 年份不存在于当前文件，创建空值
                            year_value = np.nan
                        
                        key = (year_val, indicator_col)
                        if key not in merged_data:
                            merged_data[key] = {'year': year_val, 
                                                'country_col_en': country_col_en,
                                                'country_col_ch': df.columns.get_loc('国名Ch') if '国名Ch' in df.columns else None}
                        
                        # 取当前年份和指标的数据
                        merged_data[key][indicator_col] = year_value
                
            except Exception as e:
                print(f"  - 处理文件时出错: {str(e)}")
                continue
        
        # 没有找到任何数据
        if not merged_data:
            print("警告: 未从任何文件中提取到有效数据")
            return False
        
        # 构建最终DataFrame
        final_cols = ['Year', '国名En', '国名Ch']
        final_data = {'Year': [], '国名En': [], '国名Ch': []}
        
        # 获取所有国家名称
        unique_countries_en = set()
        unique_countries_ch = set()
        
        for key, data in merged_data.items():
            country_en = df.columns.get_loc(data['country_col_en']) if data['country_col_en'] in df.columns else None
            country_ch = df.columns.get_loc('国名Ch') if '国名Ch' in df.columns else None
            
            if country_en is not None and country_ch is not None:
                unique_countries_en.add(country_en)
                unique_countries_ch.add(country_ch)
        
        # 创建所有国家组合
        country_pairs = {}
        for country_en in unique_countries_en:
            for country_ch in unique_countries_ch:
                country_pairs[(country_en, country_ch)] = True
        
        # 填充完整数据
        for (year, indicator_col), data_row in merged_data.items():
            year_val = data_row['year']
            country_en = data_row['country_col_en']
            country_ch = data_row.get('country_col_ch')
            
            # 确保国家列有效
            if country_ch is None:
                continue
            
            # 获取或创建数据行
            row_key = (year_val, country_en, country_ch)
            
            final_data['Year'].append(year_val)
            final_data['国名En'].append(country_en)
            final_data['国名Ch'].append(country_ch)
            
            # 添加指标数据
            if indicator_col not in final_cols:
                final_cols.append(indicator_col)
                final_data[indicator_col] = [np.nan] * len(final_data['Year'])
            
            idx = len(final_data['Year']) - 1
            final_data[indicator_col][idx] = merged_data.get((year_val, indicator_col), {}).get(indicator_col, np.nan)
        
        # 创建最终DataFrame
        final_df = pd.DataFrame(final_data)
        
        # 设置索引
        final_df.set_index('Year', inplace=True)
        
        # 保存到Excel文件
        try:
            with pd.ExcelWriter(output_file, engine='openpyxl', mode='w') as writer:
                final_df.to_excel(writer, sheet_name='MergedData', index=True)
            
            print(f"✅ 合并完成，成功合并 {len(indicator_files)} 个文件，共 {len(final_df)} 行数据")
            print(f"输出文件: {output_file}")
            return True
            
        except Exception as e:
            print(f"错误: 保存Excel文件时出错: {str(e)}")
            return False
            
    except Exception as e:
        print(f"合并过程中发生错误: {str(e)}")
        return False

